#ifndef __IFME_HOOKER__
#define __IFME_HOOKER__

#include <vector>

namespace ifme
{
	struct hook_items 
	{
		const char*		modname;
		const char*		funcname;
		size_t			argc;
		const void*		before_handler;
		const void*		after_handler;
		hk::hookfn_t *	hkinst;
	};

	struct hooker_t
	{
		typedef std::vector< hk::hookfn_t * > hook_instances_t;

		hook_instances_t	_hooks;
		hk::hookfn_t*		_ldr_load_dll_hkp;
		hook_items*			_hook_table;

		hooker_t( hook_items* hktab ) : _ldr_load_dll_hkp(0), _hook_table(hktab)
		{
			
		}

		long init_dll_load_notify()
		{
			_ldr_load_dll_hkp = new hk::hookfn_t(this);
			if( !_ldr_load_dll_hkp ) return E_OUTOFMEMORY;
			return _ldr_load_dll_hkp->init( LdrLoadDll, 4*4, _before_ldr_load_dll, _after_ldr_load_dll, true );
		}

		size_t before_ldr_load_dll( hk::exec_context * pctx )
		{
			PWSTR SearchPath = (PWSTR)pctx->args[0];
			PULONG DllCharacteristics = (PULONG)pctx->args[1];
			PUNICODE_STRING DllName = (PUNICODE_STRING)pctx->args[2];
			if( !DllName || !DllName->Buffer || !DllName->Length ) return true;
			HMODULE BaseNow = 0;
			NTSTATUS nts = LdrGetDllHandle( SearchPath, DllCharacteristics?*DllCharacteristics:0, DllName, &BaseNow );
			if( nts < 0 || BaseNow == NULL )
				set_thread_data( "$dllmon", xv::xvalue_t(tcs::ws_to_as(DllName->Buffer)) );
			return true;
		}
		void _after_ldr_load_dll( hk::exec_context * ctxp )
		{
			NTSTATUS nts = (NTSTATUS)pctx->retval();
			if( nts < 0 ) return ;
			PVOID * BaseAddress = (PVOID*)pctx->args[3];
			if( !BaseAddress ) return ;
			dispatch_dll_load( (HMODULE)*BaseAddress );
		}

		void dispatch_dll_load( HMODULE modbase )
		{
			for( size_t i = 0; _hook_table[i].funcname; ++ i )
			{
				hook_items & item = _hook_table[i];
				if( item.hkinst ) continue;
				hk::hookfn_t * hkp = new hk::hookfn_t();
				if( !hkp ) return ;
				HMODULE hmod = GetModuleHandleA( item.modname );
				if( !hmod ) continue;
				LPVOID fnep = GetProcAddress( hmod, item.funcname );
				if( !fnep ) continue;
				long lr = hkp->init( fnep, item.argc, item.before_handler, item.after_handler );
				if( lr < 0 ) continue;
				lr = hkp->enable();
				if( lr < 0 ) continue;
			}
		}

		static size_t _before_ldr_load_dll( hk::exec_context * ctxp )
		{
			hk::hookfn_t* self = (hk::hookfn_t*)ctxp->self;
			hooker_t * this_ = (hooker_t *)self->user_data;
			return this_->before_ldr_load_dll( ctxp );
		}

		static void _after_ldr_load_dll( hk::exec_context * ctxp )
		{
			
		}

		size_t init_protection_hooks( hook_items * items, size_t c )
		{
			size_t vc = 0;

			_hooks.reserve( c );

			for( size_t i = 0;; i < c; ++ i )
			{
				hk::hookfn_t * hkp = new hk::hookfn_t();
				if( !hkp ) break;

				// X64 ABI always ret 0
				long lr = hkp->init( items[i].target, items[i].argc * sizeof(size_t), items[i].before_handler );
				if( lr < 0 ) 
				{
					delete hkp;
					continue;
				}

				_hooks.push_back( hkp );
			}
			return vc;
		}

		size_t banch_enable()
		{
			size_t vc = 0;
			size_t c = _hooks.size();
			for( size_t i = 0; i < c; ++ i )
			{
				long lr = _hooks[i]->enable();
				if( lr < 0 ) continue;
				++ vc;
			}
			return vc;
		}

	};


};

#endif